#!/usr/bin/python

import sys
import inspect

VERSION=3

K_SRCAS                 = 1
K_DSTAS                 = 2
K_DST_PREFIX            = 3
K_DST_AS_ADDITIONAL     = 4
K_FOP                   = 5

globalArgs = {
    'format' : 1,
    'in' : [],
    'printheaders' : True,
    'printheaders_short' : False,
    'accumulateby' : K_DST_PREFIX
}

#
#   -- warn
#
def warn(x):
    sys.stderr.write("(%s) " % inspect.stack()[1][3] )
    sys.stderr.write(x)
    sys.stderr.write("\n")

def error(x):
    warn(x)
    sys.exit(1)

def print_usage_and_exit():
    print """
    --help -h 
    --format -f       1: as as 2: as as prefix
    --in -i           can be multiple times
    -p                do not print # headers
    -P                print only short # header
    -a --accumulate   'as' or 'prefix' 

    """
    print "version " , VERSION
    sys.exit(0)


#
#   -- load_options
#
def load_options():
    global globalArgs;
    argv = sys.argv[1:]
    while( len(argv) ):
        a = argv.pop(0)
        if   a in ['--help', '-h']: print_usage_and_exit()
        elif a in ['--format', '-f']:  globalArgs['format'] = int( argv.pop(0) )
        elif a in ['--in', '-i']:  globalArgs['in'].append( argv.pop(0) )
        elif a in ['-p']:  globalArgs['printheaders'] = False
        elif a in ['-P']:  globalArgs['printheaders'] = False; globalArgs['printheaders_short'] = True;
        elif a in ['-a', '--accumulate']:
            a = argv.pop(0) 
            if a == 'as': globalArgs['accumulateby'] = K_DSTAS
            elif a == 'prefix': globalArgs['accumulateby'] = K_DST_PREFIX
            else: raise Exception("unknown -a")
        else: raise Exception("unknown option %s" % a)
    return argv


def get_data_from_line_format1( line ):
    g = line.split(" ")
    r = {}
    r[K_SRCAS] = int(g[0])
    r[K_DSTAS] = int(g[1])
    r[K_FOP] = [ int(g[2]), int(g[3]), int(g[4]) ]
    return r

def get_data_from_line_format2( line ):
    g = line.split(" ")
    r = {}
    r[K_SRCAS] = int(g[0])
    r[K_DSTAS] = int(g[1])
    r[K_DST_PREFIX] = g[2]
    r[K_FOP] = [ int(g[3]), int(g[4]), int(g[5]) ]
    return r

def fop_accumulate( H, data, accumulateby):
    if accumulateby == K_DSTAS:
        k = ( data[K_SRCAS], data[K_DSTAS] )
        dst_as_additional = None
    elif accumulateby == K_DST_PREFIX:
        k = ( data[K_SRCAS], data[K_DST_PREFIX] )
        dst_as_additional = data[K_DSTAS]
    fop = data[K_FOP][:]

    if k not in H:
        H[k] = { K_FOP : fop }
        if dst_as_additional is not None:
            H[k][K_DST_AS_ADDITIONAL] = dst_as_additional

    else:
        H[k][K_FOP][0] += fop[0]
        H[k][K_FOP][1] += fop[1]
        H[k][K_FOP][2] += fop[2]

        if dst_as_additional is not None:
            if H[k][K_DST_AS_ADDITIONAL] != dst_as_additional:
                warn('DST asn conflict for prefix %s: old AS%s , new AS%s (using new)' % ( data[K_DST_PREFIX], H[k][K_DST_AS_ADDITIONAL], dst_as_additional ) )
                H[k][K_DST_AS_ADDITIONAL] = dst_as_additional
    

#
#   -- main
#
def main():
    load_options()

    H = {}
    fmt = globalArgs['format']
    files = globalArgs['in']
    printheaders = globalArgs['printheaders']
    printheaders_short = globalArgs['printheaders_short']
    accumulateby = globalArgs['accumulateby']

    if fmt == 1:
        get_data_from_line = get_data_from_line_format1
    elif fmt == 2:
        get_data_from_line = get_data_from_line_format2
    else:
        error("unknown format")

    # process files
    while(True):
        if len(files) == 0:
            o = sys.stdin
            warn("reading from stdin")
        else:
            a = files.pop(0)
            warn("opening %s" % (a) )
            o = open(a,"r")
        
        # process
        lcnt = 0
        while (True):
            lcnt = lcnt + 1
            line = o.readline()
            if not line: break

            if line[0] == "#":
                if printheaders: print line.strip("\n")
                elif printheaders_short and line[0:5] == "# in:":   print line.strip("\n")
                continue

            try:
                d = get_data_from_line(line)
            except ValueError:
                warn("    skipping: " + line.strip("\n") );
                continue

            fop_accumulate( H, d, accumulateby )

        # finalize
        warn("Read %d lines" % lcnt)
        if not o == sys.stdin: o.close()
        if len(files) == 0: break

    # display
    for k in H.keys():
        if accumulateby == K_DSTAS:
            srcas = k[0]
            dstas = k[1]
            flows = H[k][K_FOP][0]
            octets = H[k][K_FOP][1]
            packets = H[k][K_FOP][2]
            print "%s %s %s %s %s" % ( srcas, dstas, flows, octets, packets )
        elif accumulateby == K_DST_PREFIX:
            srcas = k[0]
            dstprefix = k[1]
            dstas = H[k][K_DST_AS_ADDITIONAL]
            flows = H[k][K_FOP][0]
            octets = H[k][K_FOP][1]
            packets = H[k][K_FOP][2]
            print "%s %s %s %s %s %s" % ( srcas, dstas, dstprefix, flows, octets, packets )

if __name__ == "__main__": main()


# - input format 1 or 2
# - merge by srcas-dstprefix
# - merge by srcas-dstas

